"""数据库工具模块 - 提供SQL数据库操作功能。"""

from langchain_community.utilities import SQLDatabase
from langchain_community.agent_toolkits import SQLDatabaseToolkit
from utils.models import load_chat_model


def initialize_database_tools(database_uri: str = "sqlite:///Chinook.db"):
    """Initialize database and SQL tools.
    
    Args:
        database_uri: Database connection URI
        
    Returns:
        Tuple of (database, schema_tool, query_tool, list_tables_tool)
    """
    db = SQLDatabase.from_uri(database_uri)
    llm = load_chat_model()
    toolkit = SQLDatabaseToolkit(db=db, llm=llm)
    tools = toolkit.get_tools()
    
    get_schema_tool = next(tool for tool in tools if tool.name == "sql_db_schema")
    run_query_tool = next(tool for tool in tools if tool.name == "sql_db_query")
    list_tables_tool = next(tool for tool in tools if tool.name == "sql_db_list_tables")
    
    return db, get_schema_tool, run_query_tool, list_tables_tool


# 全局工具实例
db, get_schema_tool, run_query_tool, list_tables_tool = initialize_database_tools()